import os
import time

import torch
from torch import argmax
from torch.optim import SGD
from torch.utils.data import DataLoader
from torchvision.transforms import Compose, ToTensor

from dataset import MnistDataset
from model import MLP
import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter
import datetime

if __name__ == "__main__":
    train_img_dir = "./data/MNIST/train"
    train_label_txt = "./data/MNIST/train.txt"
    test_img_dir = "./data/MNIST/test"
    test_label_txt = "./data/MNIST/test.txt"
    epoch_num = 3
    batchsize = 64
    output_root_dir = "./output"

    transform = Compose([ToTensor()])
    train_dataset = MnistDataset(img_dir_path=train_img_dir, label_txt_path=train_label_txt, transform=transform)
    test_dataset = MnistDataset(img_dir_path=test_img_dir, label_txt_path=test_label_txt, transform=transform)
    train_dl = DataLoader(train_dataset, batch_size=batchsize, shuffle=True, num_workers=4)
    test_dl = DataLoader(test_dataset, batch_size=batchsize, shuffle=False, num_workers=4)

    ## pred
    # i,(imgs,labels)=next(enumerate(train_dl))
    # print(i, imgs.shape,labels.shape)
    # model = MLP(1)
    # output = model(imgs)
    # print(output.shape)
    # pred=argmax(output,dim=1)
    # print(pred.shape)
    #
    # print(output.dtype)
    # print(labels.dtype)
    #
    # criterion=nn.CrossEntropyLoss()
    # loss=criterion(output,labels)
    # print(loss)

    ## train
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    output_dir = os.path.join(output_root_dir, timestamp)
    os.makedirs(output_dir, exist_ok=True)
    os.makedirs(os.path.join(output_dir, "logs"), exist_ok=True)
    writer = SummaryWriter(os.path.join(output_dir, "logs"))

    model = MLP(1).to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = SGD(model.parameters(), lr=0.01, momentum=0.9)
    loss_all = []


    for epoch in range(epoch_num):
        model.train()
        start_time = time.time()
        losses = 0
        for i, (imgs, labels) in enumerate(train_dl):  # 简单测试
            imgs, labels = imgs.to(device), labels.to(device)
            optimizer.zero_grad()
            output = model(imgs)
            loss = criterion(output, labels)
            losses = losses + loss.item()
            # print(loss.item())
            loss.backward()
            optimizer.step()
        losses = losses / len(train_dl)
        loss_all.append(losses)

        model_path = os.path.join(output_dir, f"model_epoch{epoch}.pth")
        torch.save(model.state_dict(), model_path)
        # print(f"model saved in {model_path}")

        pass_time = int((time.time() - start_time))  # 秒
        print(f"epoch{epoch + 1}/{epoch_num},loss:{losses}, time:{pass_time}s")
        writer.add_scalar("Loss-train", losses, epoch)


        if epoch %2 == 0:
            start_time = time.time()
            model.eval()
            val_loss = 0.0
            corract = 0
            total = 0

            with torch.no_grad():
                for imgs, labels in test_dl:
                    imgs, labels = imgs.to(device), labels.to(device)
                    output=model(imgs)
                    pred=torch.argmax(output,dim=1)
                    val_loss+=criterion(output,labels).item()*imgs.size(0)
                    corract+=(pred==labels).sum().item()
                    total+=len(labels)


            val_loss/=len(test_dl.dataset)
            corract/=total
            writer.add_scalar("Loss/val", val_loss, epoch)
            writer.add_scalar("Accuracy/val", corract, epoch)
            pass_time = int((time.time() - start_time))
            print(f" val_loss:{val_loss}, correct:{corract}, time:{pass_time}" )


    writer.close()
    model_path = os.path.join(output_dir, f"model_end.pth")
    torch.save(model.state_dict(), model_path)
    print(f"model saved in {model_path}")

    # # pred
    # model_path_now = os.path.join(output_dir, f"model_epoch{epoch}.pth")
    # model.load_state_dict(torch.load(model_path_now, map_location=device))
    # print(f"Model weights loaded from {model_path_now}")
    model.eval()
    with torch.no_grad():
        _, (imgs, labels) = next(enumerate(test_dl))
        imgs, labels = imgs.to(device), labels.to(device)
        # print(imgs.shape, labels.shape)
        output = model(imgs)
        # print(output.shape)
        pred = argmax(output, dim=1)
        # print(pred.shape)

        print(f"pred:{pred}\n, label:{labels}")

